from new_rmasm.funcs.macro_command import macro_command_translators
from new_rmasm.funcs.macro_move_command import macro_move_command_translators
from new_rmasm.translator import *
from new_rmasm.lexer import *


def jump_condition(env, ast: Union['FactorAST', 'UnaryAST', 'BinaryAST']) -> \
        List[str]:
    """ 把 ast 转换成原生语句的条件表达式
    """
    if ast is None:
        return ['always']
    if isinstance(ast, FactorAST):
        return ['equal', ast.mlog_str(env), 'true']
    if isinstance(ast, UnaryAST):
        if ast.token.code != 'not':
            raise RTranslateError(ast.token.source, ast.token.line_num,
                                  "条件运算符错误")
        return ['notEqual', ast.value.mlog_str(env), 'true']
    if isinstance(ast, BinaryAST):
        if ast.token.code not in kw.jump_ops:
            raise RTranslateError(ast.token.source, ast.token.line_num,
                                  "条件运算符错误")
        op = kw.all_operators[ast.token.code]
        return [op, ast.left.mlog_str(env), ast.right.mlog_str(env)]


def in_build_func_result_to_list(result):
    if isinstance(result, str):
        return [result]
    if isinstance(result, Iterator):
        return list(result)
    return str(result)


# 语法树根节点
class AST:
    def __init__(self, token: LexToken):
        self.token = token
        if self.token is None:
            raise Exception("NONONO")

    def translate(self, env: AstsTranslateEnv) -> List[TempCode]:
        raise self.error(f"{self.__class__} 不能直接翻译。")

    def error(self, message: str):
        return RTranslateError(self.token.source, self.token.line_num, message)

    def str(self) -> str:
        pass

    def nul_tc(self) -> TempCode:
        return TempCode(self.token, f'# ERROR Translate From <{self.str()}>')

    def __str__(self) -> str:
        # return self.str()
        return f'{self.__class__.__name__.ljust(20)} >|  {self.str()}'


# 最小数据表示，包含数字、字符串、id、raw
class FactorAST(AST):

    def str(self) -> str:
        return self.token.code

    def mlog_str(self, env: AstsTranslateEnv) -> str:
        raise self.error(f"{self.__class__} 不能获取 Mlog 数值。")


class NumberAST(FactorAST):

    def __init__(self, token: LexToken):
        super().__init__(token)
        self.negative: bool = False

    def mlog_str(self, env: AstsTranslateEnv) -> str:
        return str(self.token.get_value() * (-1 if self.negative else 1))


class StringAST(FactorAST):

    def mlog_str(self, env: AstsTranslateEnv) -> str:
        return '"' + self.token.get_value() + '"'

    def get_string_content(self) -> str:
        return self.token.get_value()

    def translate(self, env: AstsTranslateEnv) -> List[TempCode]:
        return [TempCode(self.token, self.get_string_content())]


class IdAST(FactorAST):
    def __init__(self, token: LexToken, parent: Union['IdAST', None] = None):
        super().__init__(token)
        self.parent = parent
        self.child: Union['IdAST', None] = None
        if self.parent:
            self.parent.child = self

    def set_parent(self, parent: 'IdAST'):
        self.parent = parent
        self.parent.child = self

    def str(self) -> str:
        res = f'{self.token.code}'
        if self.child:
            res = res + '.' + self.child.str()
        return res

    def mlog_str(self, env: AstsTranslateEnv) -> str:
        if self.token.code in ['result', 'self']:
            if not env.has_temp_var(self.token.code):
                raise RKeyVarNotFound(self.token.source,
                                      self.token.line_num,
                                      self.token.code)
        return env.get_var(self.token.code)

    def get_obj_and_name(self) -> List[Union['IdAST', None]]:
        nl = self.get_name_link()

        if len(nl) > 2:
            raise self.error("过多的 . 连接。")
        if len(nl) == 1:
            return [None, nl[0]]
        if len(nl) == 2:
            return [nl[0], nl[1]]

    def get_name_link(self) -> List['IdAST']:
        res: List[IdAST] = []
        start = self
        while True:
            res.append(start)
            if not start.child:
                break
            start = start.child
        return res


class RawFactorAST(FactorAST):

    def __init__(self, token: LexToken, content: List[LexToken]):
        super().__init__(token)
        self.content = content

    def mlog_str(self, env: AstsTranslateEnv) -> Union[TCToken]:
        if len(self.content) == 0:
            raise self.error("空的宏替换内容。")
        start = self.content[0]
        content = self.content[1:]
        content_len = len(content)
        if start.code == 'eval':
            return EvalTCToken(start,
                               ' '.join([i.code for i in self.content[1:]]))
        if start.code == 'position':
            if content_len == 2:
                return TagNumTCToken(start, content[0].code, content[1].code)
            elif content_len == 1:
                return TagNumTCToken(start, '', content[0].code)
            else:
                return TagNumTCToken(start, '', '')
                raise self.error("错误的 position 宏参数长度。")


# 调用方法节点，可能会单独组成语句，可能会包含在二元运算节点中（赋值调用）
class FuncAST(AST):

    def __init__(self, token: LexToken, id_ast: IdAST, args: List[FactorAST]):
        super().__init__(token)
        self.name = id_ast
        self.args = args

    def str(self) -> str:
        return f"{self.name}({', '.join([i.str() for i in self.args])})"

    def call_as_in_build(self, env: AstsTranslateEnv, result: IdAST = None) -> \
            List[TempCode]:
        obj, name = self.name.get_obj_and_name()
        obj: IdAST
        if obj:
            obj: str = obj.mlog_str(env)
        else:
            obj: None = None
        name: IdAST
        args = [a.mlog_str(env) for a in self.args]

        if result and name.str() in macro_move_command_translators:
            fun = macro_move_command_translators[name.str()]
            return [
                TempCode(self.token,
                         *in_build_func_result_to_list(
                             fun(result.mlog_str(env), obj, args)))
            ]
        if name.str() in macro_command_translators:
            fun = macro_command_translators[name.str()]
            tk = TempCode(self.token,
                          *in_build_func_result_to_list(
                              fun(obj, args)))
            if name.str() == 'comment':
                tk.step_line = False
            return [tk]
        raise self.error(f"调用内建函数 {name.str()} 失败。")

    def call_as_macro(self, env: AstsTranslateEnv, result: IdAST = None) -> List[
        TempCode]:
        obj, name = self.name.get_obj_and_name()
        obj: IdAST
        name: IdAST
        sub_env = env.new_sub_env()
        if not env.has_macro(name.str(), len(self.args)):
            raise self.error(f"不存在宏 {name.str()} 或宏不支持 {len(self.args)} 个参数。")
        macro = env.get_macro(name.str(), len(self.args))

        if result:
            sub_env.vars['result'] = result.mlog_str(env)
        if obj:
            sub_env.vars['self'] = obj.mlog_str(env)
        try:
            return macro.call(self.args, sub_env)
        except RKeyVarNotFound as e:
            if e.key == 'result':
                raise self.error(f"调用 {name.str()} 需要接收返回值。")
            if e.key == 'self':
                raise self.error(f"调用 {name.str()} 需要一个前缀对象。")

    def has_macro(self, env: AstsTranslateEnv) -> bool:
        return env.has_macro(self.name.get_obj_and_name()[1].str(),
                             len(self.args))

    def has_in_build(self, result: bool) -> bool:
        name = self.name.get_obj_and_name()[1].str()
        if result and name in macro_move_command_translators:
            return True
        return name in macro_command_translators

    def translate(self, env: AstsTranslateEnv) -> List[TempCode]:
        if self.has_in_build(False):
            return self.call_as_in_build(env)
        if self.has_macro(env):
            return self.call_as_macro(env)
        raise self.error(f"不存在 {len(self.args)} 个参数的函数 {self.name.str()}")


# 二元运算，可能会单独组成赋值或嵌套组成运算节点，也表示条件。
class BinaryAST(AST):
    def __init__(self, token: LexToken, left: FactorAST,
                 right: Union[FactorAST, 'BinaryAST']):
        super().__init__(token)
        self.right = right
        self.left = left

    def str(self) -> str:
        return f"{self.left.str()} {self.token.code} ({self.right.str()})"

    def translate(self, env: AstsTranslateEnv) -> List[TempCode]:

        if self.token.code not in kw.set_operators:
            raise self.error(f"{self.token.code} 运算符的表达式不能直接翻译。")
        if not isinstance(self.left, IdAST):
            raise self.error(f"赋值语句的左值不能是 {self.left.str()}")
        if self.token.match(LexTokenType.symbol, '='):
            if isinstance(self.right, BinaryAST):
                # 二元运算
                if self.right.token.code not in kw.all_operators:
                    raise self.error(f"未知的运算符：{self.right.token.code}")
                op = kw.all_operators[self.right.token.code]
                return [
                    TempCode(self.token,
                             'op', op, self.left.mlog_str(env),
                             self.right.left.mlog_str(env),
                             self.right.right.mlog_str(env)
                             )
                ]
            elif isinstance(self.right, UnaryAST):
                # 一元运算
                if self.right.token.code not in kw.all_operators:
                    raise self.error(f"未知的运算符：{self.right.token.code}")
                op = kw.all_operators[self.right.token.code]
                return [
                    TempCode(self.token,
                             'op', op,
                             self.left.mlog_str(env),
                             self.right.value.mlog_str(env)
                             )
                ]
            elif isinstance(self.right, FuncAST):
                # 调用函数
                func = self.right
                if func.has_in_build(True):
                    return func.call_as_in_build(env, self.left)
                if func.has_macro(env):
                    return func.call_as_macro(env, self.left)
                raise self.error(
                    f"不存在 {len(self.right.args)} 个参数的函数 {self.right.name.str()}")
            elif isinstance(self.right, FactorAST):
                # set
                return [
                    TempCode(self.token,
                             'set', self.left.mlog_str(env),
                             self.right.mlog_str(env))
                ]
            else:
                raise self.error(f"{self.token.code} 表达式右侧翻译错误。")
        elif self.token.code in kw.set_operators:
            if not isinstance(self.right, FactorAST):
                raise self.error(f"{self.token.code} 表达式的右侧必须是一个数值，不能是表达式。")
            # 二元运算
            op = self.token.code.replace('=', '')
            if op not in kw.all_operators:
                raise self.error(f"未知的运算符：{self.token.code}")
            op = kw.all_operators[op]
            return [
                TempCode(self.token, 'op', op,
                         self.left.mlog_str(env),
                         self.left.mlog_str(env),
                         self.right.mlog_str(env)
                         )
            ]


# 一元运算，可能组成运算节点的值部分，也表示条件。
class UnaryAST(AST):

    def __init__(self, token: LexToken, value: FactorAST):
        super().__init__(token)
        self.value = value

    def str(self) -> str:
        return f"{self.token.code} ({self.value.str()})"

    def is_negative(self) -> bool:
        return self.token.match(LexTokenType.symbol, '-') and \
               isinstance(self.value, FactorAST)


# 标签
class TagAST(AST):

    def __init__(self, token: LexToken, name: str):
        super().__init__(token)
        self.name = name

    def str(self) -> str:
        return f"{self.name}{self.token.code}"

    def translate(self, env: AstsTranslateEnv) -> List[TempCode]:
        tc = TempCode(self.token, TagTCToken(self.token, self.name))
        tc.has_result = False
        tc.step_line = False
        return [tc]


# 跳转
class JumpAST(AST):

    def __init__(self, token: LexToken, direction: str, target: str,
                 condition: Union[BinaryAST, UnaryAST, FactorAST]):
        super().__init__(token)
        self.direction = direction
        self.condition = condition
        self.target = target

    def str(self) -> str:
        res = f'jump {self.direction}{self.target}'
        if self.condition:
            res += f' if {self.condition.str()}'
        return res

    def translate(self, env: AstsTranslateEnv) -> List[TempCode]:
        if self.condition:
            cond = jump_condition(env, self.condition)
        else:
            cond = ['always']
        tc = TempCode(
            self.token, 'jump',
            TagNumTCToken(self.token, self.direction, self.target),
            *cond
        )

        return [tc]


# end
class EndAST(AST):
    def translate(self, env: AstsTranslateEnv) -> List[TempCode]:
        return [TempCode(self.token, 'end')]


# {} 包裹的一段代码，也用此来解析一个文件。
class SegmentAST(AST):

    def __init__(self, token: Union[None, LexToken]):
        super().__init__(token)
        self.asts: List[AST] = []

    def str(self) -> str:
        if not self.asts:
            return '{ }'
        res = []
        for a in self.asts:
            res.append(str(a))

        return '{\n' + '\n'.join(res) + '\n}'

    def translate(self, env: AstsTranslateEnv, sub_env: bool = True) -> List[
        TempCode]:
        res = []
        if sub_env:
            env = env.new_sub_env()
        for a in self.asts:
            r = a.translate(env)
            res.extend(r)
        return res


# define
class DefineMacroAST(AST):
    def __init__(self, token: LexToken, name: str, args: List[FactorAST],
                 body: SegmentAST):
        super().__init__(token)
        self.args = args
        self.name = name
        self.body = body

    def str(self) -> str:
        if self.body.asts:
            body = '{' + self.body.str() + '\n}'
        else:
            body = '{ }'
        args = f"({', '.join([i.str() for i in self.args])})"
        return f"define {self.name} {args} {body}"

    def translate(self, env: AstsTranslateEnv) -> List[TempCode]:
        env.new_macro(self)
        return []

    def call(self, args: List, call_in_env: AstsTranslateEnv) -> List[TempCode]:
        if len(args) != len(self.args):
            raise self.error(
                f"{self.name} 需要 {len(self.args)} 个参数，而不是 {len(args)} 个。")

        for k, v in enumerate(args):
            if isinstance(v, SegmentAST):
                call_in_env.macros[(self.args[k].str(), 0)] = DefineMacroAST(
                    v.token, self.args[k].str(), [], v
                )
            elif isinstance(v, DefineMacroAST):
                call_in_env.macros[(self.args[k].str(), len(v.args))] = v
            else:
                call_in_env.vars[self.args[k].str()] = v.mlog_str(call_in_env)
        return self.body.translate(call_in_env, False)


# let
class DefineVarAST(AST):
    def __init__(self, token: LexToken,
                 var_list: List[Union[IdAST, BinaryAST]]):
        super().__init__(token)
        self.vars = var_list

    def str(self) -> str:
        vs = []
        for v in self.vars:
            if isinstance(v, FactorAST):
                vs.append(v.str())
            if isinstance(v, BinaryAST):
                vs.append(f"{v.left.str()} {v.token.code} {v.right.str()}")
        return f"{self.token.code} {', '.join(vs)}"


# let
class DefineTempVarAST(DefineVarAST):
    def translate(self, env: AstsTranslateEnv) -> List[TempCode]:
        res: List[TempCode] = []
        for v in self.vars:
            if isinstance(v, IdAST):
                if v.str() in env.vars:
                    raise self.error(f"重复的临时变量：{v.str()}")
                env.new_temp_var(v.str())
                env.add_war(self.error(f"小心临时变量 {v.str()} 没有默认值。"))
            elif isinstance(v, BinaryAST) and \
                    v.token.code == '=' and \
                    isinstance(v.left, IdAST):
                if v.left.str() in env.vars:
                    raise self.error(f"重复的临时变量：{v.left.str()}")
                env.new_temp_var(v.left.str())
                res.extend(v.translate(env))
            else:
                raise self.error("变量声明语句错误。")
        return res


# if  ... else
class IfAST(AST):

    def __init__(self, token: LexToken,
                 condition: Union[BinaryAST, FactorAST, UnaryAST],
                 then_body: SegmentAST,
                 else_body: Union[SegmentAST, 'IfAST'] = None):
        super().__init__(token)
        self.else_body = else_body
        self.then_body = then_body
        self.condition = condition

    def str(self) -> str:
        res = "if " + self.condition.str() + ' ' + self.then_body.str() + '\n'
        if isinstance(self.else_body, SegmentAST):
            res += ' else ' + self.else_body.str() + '\n'
        elif isinstance(self.else_body, IfAST):
            res += ' else ' + self.else_body.str()
        return res

    def translate(self, env: AstsTranslateEnv) -> List[TempCode]:
        else_tcs = []
        then_tag = env.new_random_tag()
        end_tag = env.new_random_tag()
        if self.else_body:
            else_tcs = self.else_body.translate(env.new_sub_env())
        res: List[TempCode] = [
            TempCode(self.token,
                     'jump', TagNumTCToken(self.token, '+', then_tag),
                     *jump_condition(env, self.condition)
                     ),
            *else_tcs,
            TempCode(self.token,
                     'jump', TagNumTCToken(self.token, '+', end_tag)
                     ),
            TagTempCode(self.token, then_tag),
            *self.then_body.translate(env.new_sub_env()),
            TagTempCode(self.token, end_tag)
        ]
        return res


# for ;;
class ForAST(AST):

    def __init__(self, token: LexToken, init: AST,
                 condition: Union[BinaryAST, UnaryAST, FactorAST],
                 loop_end: AST, body: SegmentAST):
        super().__init__(token)
        self.body = body
        self.loop_end = loop_end
        self.condition = condition
        self.init = init

    def str(self) -> str:
        res = f'for {self.init.str()}; {self.condition.str()}; {self.loop_end.str()} '
        res += '' + self.body.str() + ''
        return res

    def translate(self, env: AstsTranslateEnv) -> List[TempCode]:
        env = env.new_sub_env()
        res: List[TempCode] = [*self.init.translate(env)]
        loop_check_tag = env.new_random_tag()
        loop_start_tag = env.new_random_tag()
        loop_end_tag = env.new_random_tag()
        res.extend([
            TagTempCode(self.token, loop_check_tag),
            TempCode(self.token, 'jump',
                     TagNumTCToken(self.token, '+', loop_start_tag),
                     *jump_condition(env, self.condition)),
            TempCode(self.token, 'jump',
                     TagNumTCToken(self.token, '+', loop_end_tag)),
            TagTempCode(self.token, loop_start_tag),
            *self.body.translate(env),
            TempCode(self.token, 'jump',
                     TagNumTCToken(self.token, '-', loop_check_tag),
                     *jump_condition(env, self.condition)),
            TagTempCode(self.token, loop_end_tag),
        ])

        return res


# while
class WhileAST(AST):

    def __init__(self, token: LexToken,
                 condition: Union[BinaryAST, UnaryAST, FactorAST],
                 body: SegmentAST):
        super().__init__(token)
        self.condition = condition
        self.body = body

    def str(self) -> str:
        res = f'while {self.condition.str()}'
        res += '' + self.body.str() + ''
        return res

    def translate(self, env: AstsTranslateEnv) -> List[TempCode]:
        end_tag = env.new_random_tag()
        loop_tag = env.new_random_tag()

        res: List[TempCode] = [
            TempCode(self.token, 'jump',
                     TagNumTCToken(self.token, '+', loop_tag),
                     *jump_condition(env, self.condition)),
            TempCode(self.token, 'jump',
                     TagNumTCToken(self.token, '+', end_tag)),
            TagTempCode(self.token, loop_tag),
            *self.body.translate(env),
            TempCode(self.token, 'jump',
                     TagNumTCToken(self.token, '-', loop_tag),
                     *jump_condition(env, self.condition)),
            TagTempCode(self.token, end_tag)
        ]
        return res


# import "xx"
class ImportAST(AST):

    def __init__(self, token: LexToken, file_path: StringAST):
        super().__init__(token)
        self.file_path = file_path

    def str(self) -> str:
        return f"import {self.file_path.str()}"

    def translate(self, env: AstsTranslateEnv) -> List[TempCode]:
        file_path = self.file_path.get_string_content()
        if env.has_imported(file_path):
            # env.add_war(self.error(
            #     f"重复的 import 会被忽略：{file_path}"))
            return []
        from new_rmasm import lexer, parser
        code = utils.open_code_file(file_path)
        lex = lexer.Lexer(code, source=self.file_path.get_string_content())
        segment = parser.AstParser(lex).segment_ast()

        sub = env.new_sub_env()
        sub.macros = env.macros
        res = segment.translate(sub, False)
        env.add_imported(file_path)
        return res
